"""Agniv's code - do not distribute before repo release - modified for bf16 and mixed precision
has a bunch of NaNs for now, maybe will fix later if desperate
"""

import math
import torch

import triton
import triton.language as tl

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    import torch.version


def is_hip():
    return torch.version.hip


@triton.jit
def _attn_fwd_inner(
    acc,
    l_i,
    m_i,
    q,  #
    K_block_ptr,
    V_block_ptr,
    Mask_block_ptr,  #
    blkBinaryMask_ptr,
    num_ones,
    offset_ones,  #
    start_m,
    qk_scale,  #
    blkBinaryMask_stride_d1,
    blkBinaryMask_stride_d2,  # Strides for BLK_ptr
    num_ones_stride_d1,  #
    IS_DENSEBLK: tl.constexpr,  #
    BLOCK_M: tl.constexpr,
    HEAD_DIM: tl.constexpr,
    BLOCK_N: tl.constexpr,  #
    N_CTX: tl.constexpr,
    fp8_v: tl.constexpr,
):
    lo, hi = 0, N_CTX
    K_block_ptr = tl.advance(K_block_ptr, (0, lo))
    V_block_ptr = tl.advance(V_block_ptr, (lo, 0))
    Mask_block_ptr = tl.advance(Mask_block_ptr, (0, lo))

    num_ones_val = 0
    offset_ones_val = 0
    total_ones_val = 0
    if IS_DENSEBLK == 1:
        num_ones_ptr = start_m * num_ones_stride_d1
        offset_ones_ptr = start_m * num_ones_stride_d1
        num_ones_val = tl.load(num_ones + num_ones_ptr)
        offset_ones_val = tl.load(offset_ones + offset_ones_ptr)
        total_ones_val = num_ones_val + offset_ones_val

    # loop over k, v and update accumulator
    for start_n in range(lo, hi, BLOCK_N):
        cur_val = start_n // BLOCK_N
        start_n = tl.multiple_of(start_n, BLOCK_N)

        binaryBlk_ptr = start_m * blkBinaryMask_stride_d1 + cur_val * blkBinaryMask_stride_d2
        binaryBlk_val = tl.load(blkBinaryMask_ptr + binaryBlk_ptr)
        if binaryBlk_val:
            # if offset_ones_val != 0:
            #     offset_ones_val = offset_ones_val - 1
            # elif num_ones_val != 0:
            #     num_ones_val =num_ones_val - 1
            #     if MASK ==1:
            #         MASK = 0
            # elif num_ones_val ==0 and MASK == 0:
            #     MASK = 1

            # -- compute qk ----
            k = tl.load(K_block_ptr)
            qk = tl.dot(q, k)

            ## Added for random masking case
            if cur_val < offset_ones_val and IS_DENSEBLK == 1:
                mask = tl.load(Mask_block_ptr)
                qk = qk * qk_scale + tl.where(mask, 0, -1.0e6)
                m_ij = tl.maximum(m_i, tl.max(qk, 1))
                qk -= m_ij[:, None]
            elif cur_val < total_ones_val and IS_DENSEBLK == 1:
                qk = qk * qk_scale
                m_ij = tl.maximum(m_i, tl.max(qk, 1))
                qk -= m_ij[:, None]
                # print('num_ones_val', num_ones_val)
            else:
                mask = tl.load(Mask_block_ptr)
                qk = qk * qk_scale + tl.where(mask, 0, -1.0e6)
                m_ij = tl.maximum(m_i, tl.max(qk, 1))
                qk -= m_ij[:, None]

            p = tl.math.exp2(qk)
            l_ij = tl.sum(p, 1)
            # -- update m_i and l_i
            alpha = tl.math.exp2(m_i - m_ij)
            l_i = l_i * alpha + l_ij
            # -- update output accumulator --
            acc = acc * alpha[:, None]
            # update acc
            v = tl.load(V_block_ptr)
            p = p.to(tl.float8e5) if fp8_v else p.to(q.type.element_ty)
            acc = tl.dot(p, v, acc)
            # update m_i and l_i
            m_i = m_ij
        V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0))
        K_block_ptr = tl.advance(K_block_ptr, (0, BLOCK_N))
        Mask_block_ptr = tl.advance(Mask_block_ptr, (0, BLOCK_N))
        # blkBinaryMask_ptr = tl.advance(blkBinaryMask_ptr, (0, 1))
    return acc, l_i, m_i


@triton.jit
def _attn_fwd(
    Q,
    K,
    V,
    sm_scale,
    M,
    Out,
    maskMat,
    blkBinaryMask,  #
    # num_ones tells number of blocks where all entries are ones, and
    # offset_ones tells the offset of the first such block.
    # this setting is used for causal/dense attention.
    # num_ones_stride_d1 adn d2 is common stride for num_ones and offset_ones.
    num_ones,
    offset_ones,
    num_ones_stride_d1,  #
    stride_qz,
    stride_qh,
    stride_qm,
    stride_qk,  #
    stride_kz,
    stride_kh,
    stride_kn,
    stride_kk,  #
    stride_vz,
    stride_vh,
    stride_vk,
    stride_vn,  #
    stride_oz,
    stride_oh,
    stride_om,
    stride_on,  #
    maskMat_stride_d1,
    maskMat_stride_d2,  #
    blkBinaryMask_stride_d1,
    blkBinaryMask_stride_d2,  #
    Z,
    H,
    N_CTX,  #
    IS_DENSEBLK: tl.constexpr,  #
    HEAD_DIM: tl.constexpr,  #
    BLOCK_M: tl.constexpr,  #
    BLOCK_N: tl.constexpr,  #
):
    tl.static_assert(BLOCK_N <= HEAD_DIM)
    start_m = tl.program_id(0)
    off_hz = tl.program_id(1)
    off_z = off_hz // H
    off_h = off_hz % H
    qvk_offset = off_z.to(tl.int64) * stride_qz + off_h.to(tl.int64) * stride_qh

    # block pointers
    Q_block_ptr = tl.make_block_ptr(
        base=Q + qvk_offset,
        shape=(N_CTX, HEAD_DIM),
        strides=(stride_qm, stride_qk),
        offsets=(start_m * BLOCK_M, 0),
        block_shape=(BLOCK_M, HEAD_DIM),
        order=(1, 0),
    )
    v_order: tl.constexpr = (0, 1) if V.dtype.element_ty == tl.float8e5 else (1, 0)  # type: ignore
    V_block_ptr = tl.make_block_ptr(
        base=V + qvk_offset,
        shape=(N_CTX, HEAD_DIM),
        strides=(stride_vk, stride_vn),
        offsets=(0, 0),
        block_shape=(BLOCK_N, HEAD_DIM),
        order=v_order,
    )
    K_block_ptr = tl.make_block_ptr(
        base=K + qvk_offset,
        shape=(HEAD_DIM, N_CTX),
        strides=(stride_kk, stride_kn),
        offsets=(0, 0),
        block_shape=(HEAD_DIM, BLOCK_N),
        order=(0, 1),
    )
    O_block_ptr = tl.make_block_ptr(
        base=Out + qvk_offset,
        shape=(N_CTX, HEAD_DIM),
        strides=(stride_om, stride_on),
        offsets=(start_m * BLOCK_M, 0),
        block_shape=(BLOCK_M, HEAD_DIM),
        order=(1, 0),
    )
    Mask_block_ptr = tl.make_block_ptr(
        base=maskMat,
        shape=(N_CTX, N_CTX),
        strides=(maskMat_stride_d1, maskMat_stride_d2),
        offsets=(start_m * BLOCK_M, 0),
        block_shape=(BLOCK_M, BLOCK_N),
        order=(1, 0),
    )

    # initialize offsets
    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
    # offs_n = tl.arange(0, BLOCK_N)

    # initialize pointer to m and l
    m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
    l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0
    acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32)
    # load scales
    qk_scale = sm_scale
    qk_scale *= 1.44269504  # 1/log(2)
    # load q: it will stay in SRAM throughout
    q = tl.load(Q_block_ptr)

    acc, l_i, m_i = _attn_fwd_inner(
        acc,
        l_i,
        m_i,
        q,
        K_block_ptr,
        V_block_ptr,
        Mask_block_ptr,  #
        blkBinaryMask,
        num_ones,
        offset_ones,  #
        start_m,
        qk_scale,  #
        blkBinaryMask_stride_d1,
        blkBinaryMask_stride_d2,  #
        num_ones_stride_d1,
        IS_DENSEBLK,
        BLOCK_M,
        HEAD_DIM,
        BLOCK_N,  #
        N_CTX,
        V.dtype.element_ty == tl.float8e5,  #
    )  # type: ignore

    # epilogue
    m_i += tl.math.log2(l_i)
    acc = acc / l_i[:, None]
    m_ptrs = M + off_hz * N_CTX + offs_m
    tl.store(m_ptrs, m_i)
    tl.store(O_block_ptr, acc.to(Out.type.element_ty))


@triton.jit
def _attn_bwd_preprocess(
    O,
    DO,  #
    Delta,  #
    Z,
    H,
    N_CTX,  #
    BLOCK_M: tl.constexpr,
    HEAD_DIM: tl.constexpr,  #
):
    off_m = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
    off_hz = tl.program_id(1)
    off_n = tl.arange(0, HEAD_DIM)
    # load
    o = tl.load(O + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :])
    do = tl.load(DO + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :]).to(tl.float32)
    delta = tl.sum(o * do, axis=1)
    # write-back
    tl.store(Delta + off_hz * N_CTX + off_m, delta)


# The main inner-loop logic for computing dK and dV.
@triton.jit
def _attn_bwd_dkdv(
    dk,
    dv,  #
    Q,
    k,
    v,
    sm_scale,
    maskMat,
    blkBinaryMask,  #
    num_ones,
    offset_ones,  #
    DO,  #
    M,
    D,  #
    # shared by Q/K/V/DO.
    stride_tok,
    stride_d,  #
    stride_m_d1,
    stride_m_d2,  #
    blkBinaryMask_stride_d1,
    blkBinaryMask_stride_d2,  #
    num_ones_stride_d1,  #
    H,
    N_CTX,
    IS_DENSEBLK: tl.constexpr,  #
    BLOCK_M1: tl.constexpr,  #
    BLOCK_N1: tl.constexpr,  #
    HEAD_DIM: tl.constexpr,  #
    # Filled in by the wrapper.
    start_n,
    start_m,
    num_steps,  #
):
    offs_m = start_m + tl.arange(0, BLOCK_M1)
    offs_n = start_n + tl.arange(0, BLOCK_N1)
    offs_k = tl.arange(0, HEAD_DIM)
    qT_ptrs = Q + offs_m[None, :] * stride_tok + offs_k[:, None] * stride_d
    do_ptrs = DO + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
    maskMat_T_ptrs = maskMat + offs_m[None, :] * stride_m_d1 + offs_n[:, None] * stride_m_d2
    curr_n = start_n // BLOCK_N1
    num_ones_val = 0
    offset_ones_val = 0
    total_ones_val = 0
    if IS_DENSEBLK == 1:
        num_ones_ptr = curr_n * num_ones_stride_d1
        offset_ones_ptr = curr_n * num_ones_stride_d1
        num_ones_val = tl.load(num_ones + num_ones_ptr)
        offset_ones_val = tl.load(offset_ones + offset_ones_ptr)
        total_ones_val = num_ones_val + offset_ones_val

    # BLOCK_N1 must be a multiple of BLOCK_M1, otherwise the code wouldn't work.
    tl.static_assert(BLOCK_N1 % BLOCK_M1 == 0)
    curr_m = start_m
    step_m = BLOCK_M1
    for blk_idx in range(num_steps):
        blkBinaryMask_ptr = curr_n * blkBinaryMask_stride_d1 + blk_idx * blkBinaryMask_stride_d2
        binaryBlk_val = tl.load(blkBinaryMask + blkBinaryMask_ptr)
        if binaryBlk_val:
            qT = tl.load(qT_ptrs)
            # Load m before computing qk to reduce pipeline stall.
            offs_m = curr_m + tl.arange(0, BLOCK_M1)
            m = tl.load(M + offs_m)
            qkT = tl.dot(k, qT)
            pT = tl.math.exp2(qkT - m[None, :])
            # Autoregressive masking.
            if blk_idx < offset_ones_val and IS_DENSEBLK == 1:
                mask = tl.load(maskMat_T_ptrs)
                pT = tl.where(mask, pT, 0.0)
            elif blk_idx < total_ones_val and IS_DENSEBLK == 1:
                pT = pT
            else:
                mask = tl.load(maskMat_T_ptrs)
                pT = tl.where(mask, pT, 0.0)
            do = tl.load(do_ptrs)
            # Compute dV.
            ppT = pT
            ppT = ppT.to(Q.type.element_ty)
            dv += tl.dot(ppT, do)
            # D (= delta) is pre-divided by ds_scale.
            Di = tl.load(D + offs_m)
            # Compute dP and dS.
            dpT = tl.dot(v, tl.trans(do)).to(tl.float32)
            dsT = pT * (dpT - Di[None, :])
            dsT = dsT.to(Q.type.element_ty)
            dk += tl.dot(dsT, tl.trans(qT))
        # Increment pointers.
        curr_m += step_m
        qT_ptrs += step_m * stride_tok
        do_ptrs += step_m * stride_tok
        maskMat_T_ptrs += step_m * stride_m_d1
    return dk, dv


# the main inner-loop logic for computing dQ
@triton.jit
def _attn_bwd_dq(
    dq,
    q,
    K,
    V,
    maskMat,
    blkBinaryMask,  #
    num_ones,
    offset_ones,  #
    do,
    m,
    D,
    # shared by Q/K/V/DO.
    stride_tok,
    stride_d,  #
    stride_m_d1,
    stride_m_d2,  #
    blkBinaryMask_stride_d1,
    blkBinaryMask_stride_d2,  #
    num_ones_stride_d1,  #
    H,
    N_CTX,  #
    IS_DENSEBLK: tl.constexpr,  #
    BLOCK_M2: tl.constexpr,  #
    BLOCK_N2: tl.constexpr,  #
    HEAD_DIM: tl.constexpr,
    # Filled in by the wrapper.
    start_m,
    start_n,
    num_steps,  #
):
    offs_m = start_m + tl.arange(0, BLOCK_M2)
    offs_n = start_n + tl.arange(0, BLOCK_N2)
    offs_k = tl.arange(0, HEAD_DIM)
    kT_ptrs = K + offs_n[None, :] * stride_tok + offs_k[:, None] * stride_d
    vT_ptrs = V + offs_n[None, :] * stride_tok + offs_k[:, None] * stride_d
    maskMat_ptrs = maskMat + offs_m[:, None] * stride_m_d1 + offs_n[None, :] * stride_m_d2
    curr_m = start_m // BLOCK_M2
    ## Changes for checking dense blocks
    num_ones_val = 0
    offset_ones_val = 0
    total_ones_val = 0
    if IS_DENSEBLK == 1:
        num_ones_ptr = curr_m * num_ones_stride_d1
        offset_ones_ptr = curr_m * num_ones_stride_d1
        num_ones_val = tl.load(num_ones + num_ones_ptr)
        offset_ones_val = tl.load(offset_ones + offset_ones_ptr)
        total_ones_val = num_ones_val + offset_ones_val

    # D (= delta) is pre-divided by ds_scale.
    Di = tl.load(D + offs_m)
    # BLOCK_M2 must be a multiple of BLOCK_N2, otherwise the code wouldn't work.
    tl.static_assert(BLOCK_M2 % BLOCK_N2 == 0)
    curr_n = start_n
    step_n = BLOCK_N2
    for blk_idx in range(num_steps):
        binaryBlk_ptr = curr_m * blkBinaryMask_stride_d1 + blk_idx * blkBinaryMask_stride_d2
        binaryBlk_val = tl.load(blkBinaryMask + binaryBlk_ptr)
        # print('blk_idx', blk_idx)
        if binaryBlk_val:
            kT = tl.load(kT_ptrs)
            vT = tl.load(vT_ptrs)
            qk = tl.dot(q, kT)
            p = tl.math.exp2(qk - m)
            # Autoregressive masking.
            if blk_idx < offset_ones_val and IS_DENSEBLK == 1:
                mask = tl.load(maskMat_ptrs)
                p = tl.where(mask, p, 0.0)
            elif blk_idx < total_ones_val and IS_DENSEBLK == 1:
                p = p
            else:
                mask = tl.load(maskMat_ptrs)
                # offs_n = curr_n + tl.arange(0, BLOCK_N2)
                # mask = (offs_m[:, None] >= offs_n[None, :])
                p = tl.where(mask, p, 0.0)
            # Compute dP and dS.
            dp = tl.dot(do, vT).to(tl.float32)
            ds = p * (dp - Di[:, None])
            ds = ds.to(q.type.element_ty)
            # Compute dQ.
            # NOTE: We need to de-scale dq in the end, because kT was pre-scaled.
            dq += tl.dot(ds, tl.trans(kT))
        # Increment pointers.
        curr_n += step_n
        kT_ptrs += step_n * stride_tok
        vT_ptrs += step_n * stride_tok
        maskMat_ptrs += step_n * stride_m_d2

    return dq


@triton.jit
def _attn_bwd(
    Q,
    K,
    V,
    sm_scale,
    maskMat,
    blkBinaryMask,
    blkBinaryMask_T,  #
    num_ones,
    offset_ones,
    num_ones_stride_d1,  #
    num_ones_T,
    offset_ones_T,
    num_ones_T_stride_d1,  #
    DO,  #
    DQ,
    DK,
    DV,  #
    M,
    D,
    # shared by Q/K/V/DO.
    stride_z,
    stride_h,
    stride_tok,
    stride_d,  #
    stride_m_d1,
    stride_m_d2,  #
    blkBinaryMask_stride_d1,
    blkBinaryMask_stride_d2,  #
    blkBinaryMask_T_stride_d1,
    blkBinaryMask_T_stride_d2,  #
    H,
    N_CTX,  #
    IS_DENSEBLK: tl.constexpr,  #
    BLOCK_M1: tl.constexpr,  #
    BLOCK_N1: tl.constexpr,  #
    BLOCK_M2: tl.constexpr,  #
    BLOCK_N2: tl.constexpr,  #
    BLK_SLICE_FACTOR: tl.constexpr,  #
    HEAD_DIM: tl.constexpr,
):
    LN2: tl.constexpr = 0.6931471824645996  # = ln(2)  # type: ignore

    bhid = tl.program_id(2)
    off_chz = (bhid * N_CTX).to(tl.int64)
    adj = (stride_h * (bhid % H) + stride_z * (bhid // H)).to(tl.int64)
    pid = tl.program_id(0)

    # offset pointers for batch/head
    Q += adj
    K += adj
    V += adj
    DO += adj
    DQ += adj
    DK += adj
    DV += adj
    M += off_chz
    D += off_chz
    ## NO MaskMat update till here

    # load scales
    offs_k = tl.arange(0, HEAD_DIM)

    start_n = pid * BLOCK_N1
    # start_m = start_n
    start_m = 0

    # MASK_BLOCK_M1: tl.constexpr = BLOCK_M1 // BLK_SLICE_FACTOR
    offs_n = start_n + tl.arange(0, BLOCK_N1)

    dv = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)
    dk = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)

    # load K and V: they stay in SRAM throughout the inner loop.
    k = tl.load(K + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)
    v = tl.load(V + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)

    # num_steps = BLOCK_N1 // MASK_BLOCK_M1

    # dk, dv = _attn_bwd_dkdv(dk, dv,  #
    #                         Q, k, v, sm_scale,  #
    #                         DO,  #
    #                         M, D,  #
    #                         stride_tok, stride_d,  #
    #                         H, N_CTX,  #
    #                         MASK_BLOCK_M1, BLOCK_N1, HEAD_DIM,  #
    #                         start_n, start_m, num_steps,  #
    #                         MASK=True  #
    #                         )

    # start_m += num_steps * MASK_BLOCK_M1
    num_steps = (N_CTX) // BLOCK_M1

    # Compute dK and dV for non-masked blocks.
    dk, dv = _attn_bwd_dkdv(  #
        dk,
        dv,  #
        Q,
        k,
        v,
        sm_scale,
        maskMat,
        blkBinaryMask_T,  #
        num_ones_T,
        offset_ones_T,  #
        DO,  #
        M,
        D,  #
        stride_tok,
        stride_d,  #
        stride_m_d1,
        stride_m_d2,
        blkBinaryMask_T_stride_d1,
        blkBinaryMask_T_stride_d2,  #
        num_ones_T_stride_d1,  #
        H,
        N_CTX,  #
        IS_DENSEBLK,  #
        BLOCK_M1,
        BLOCK_N1,
        HEAD_DIM,  #
        start_n,
        start_m,
        num_steps,
    )  # type: ignore

    dv_ptrs = DV + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
    tl.store(dv_ptrs, dv)

    # Write back dK.
    dk *= sm_scale
    dk_ptrs = DK + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
    tl.store(dk_ptrs, dk)

    # THIS BLOCK DOES DQ:
    start_m = pid * BLOCK_M2
    end_n = N_CTX

    # MASK_BLOCK_N2: tl.constexpr = BLOCK_N2 // BLK_SLICE_FACTOR
    offs_m = start_m + tl.arange(0, BLOCK_M2)

    q = tl.load(Q + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d)
    dq = tl.zeros([BLOCK_M2, HEAD_DIM], dtype=tl.float32)
    do = tl.load(DO + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d)

    m = tl.load(M + offs_m)
    m = m[:, None]

    # Compute dQ for masked (diagonal) blocks.
    # NOTE: This code scans each row of QK^T backward (from right to left,
    # but inside each call to _attn_bwd_dq, from left to right), but that's
    # not due to anything important.  I just wanted to reuse the loop
    # structure for dK & dV above as much as possible.
    # num_steps = BLOCK_M2 // MASK_BLOCK_N2
    # dq = _attn_bwd_dq(dq, q, K, V,  #
    #                   do, m, D,  #
    #                   stride_tok, stride_d,  #
    #                   H, N_CTX,  #
    #                   BLOCK_M2, MASK_BLOCK_N2, HEAD_DIM,  #
    #                   start_m, end_n - num_steps * MASK_BLOCK_N2, num_steps,  #
    #                   MASK=True  #
    #                   )
    # end_n -= num_steps * MASK_BLOCK_N2
    # # stage 2

    num_steps = end_n // BLOCK_N2
    dq = _attn_bwd_dq(
        dq,
        q,
        K,
        V,
        maskMat,
        blkBinaryMask,  #
        num_ones,
        offset_ones,  #
        do,
        m,
        D,  #
        stride_tok,
        stride_d,  #
        stride_m_d1,
        stride_m_d2,  #
        blkBinaryMask_stride_d1,
        blkBinaryMask_stride_d2,  #
        num_ones_stride_d1,  #
        H,
        N_CTX,  #
        IS_DENSEBLK,  #
        BLOCK_M2,
        BLOCK_N2,
        HEAD_DIM,  #
        start_m,
        end_n - num_steps * BLOCK_N2,
        num_steps,
    )
    # Write back dQ.
    dq_ptrs = DQ + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
    dq *= LN2
    tl.store(dq_ptrs, dq)


class _attention_dense_binaryBlkMat(torch.autograd.Function):
    """
    This is a cleaned up version of the Binary Block Masked Attention.
    It doesn't contain the casual masking logic.
    """

    @staticmethod
    def forward(
        ctx,
        q,
        k,
        v,
        causal,
        sm_scale,
        maskMat,
        blkBinaryMask,
        blkBinaryMask_T,
        num_ones,
        offset_ones,
        num_ones_T,
        offset_ones_T,
    ):
        # shape constraints
        HEAD_DIM_Q, HEAD_DIM_K = q.shape[-1], k.shape[-1]
        # when v is in float8_e5m2 it is transposed.
        HEAD_DIM_V = v.shape[-1]
        assert HEAD_DIM_Q == HEAD_DIM_K and HEAD_DIM_K == HEAD_DIM_V
        assert HEAD_DIM_K in {16, 32, 64, 128, 256}
        o = torch.empty_like(q)
        stage = 1 if causal else 0
        extra_kern_args = {}
        # Tuning for AMD target
        if is_hip():
            waves_per_eu = 3 if HEAD_DIM_K <= 64 else 2
            extra_kern_args = {"waves_per_eu": waves_per_eu, "allow_flush_denorm": True}

        grid = lambda args: (triton.cdiv(q.shape[2], args["BLOCK_M"]), q.shape[0] * q.shape[1], 1)
        M = torch.empty((q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32)
        _attn_fwd[grid](
            q,
            k,
            v,
            sm_scale,
            M,
            o,
            maskMat,
            blkBinaryMask,  #
            num_ones,
            offset_ones,
            num_ones.stride(0),  #
            q.stride(0),
            q.stride(1),
            q.stride(2),
            q.stride(3),  #
            k.stride(0),
            k.stride(1),
            k.stride(2),
            k.stride(3),  #
            v.stride(0),
            v.stride(1),
            v.stride(2),
            v.stride(3),  #
            o.stride(0),
            o.stride(1),
            o.stride(2),
            o.stride(3),  #
            maskMat.stride(0),
            maskMat.stride(1),  #
            blkBinaryMask.stride(0),
            blkBinaryMask.stride(1),  #
            q.shape[0],
            q.shape[1],  #
            N_CTX=q.shape[2],  #
            IS_DENSEBLK=stage,  # type: ignore
            HEAD_DIM=HEAD_DIM_K,  #
            BLOCK_M=128,  # type: ignore
            BLOCK_N=32,  # type: ignore
            **extra_kern_args,
        )

        ctx.save_for_backward(
            q, k, v, o, M, maskMat, blkBinaryMask, blkBinaryMask_T, num_ones, offset_ones, num_ones_T, offset_ones_T
        )
        ctx.grid = grid
        ctx.sm_scale = sm_scale
        ctx.HEAD_DIM = HEAD_DIM_K
        ctx.causal = causal
        return o

    @staticmethod
    def backward(ctx, do):
        q, k, v, o, M, maskMat, blkBinaryMask, blkBinaryMask_T, num_ones, offset_ones, num_ones_T, offset_ones_T = (
            ctx.saved_tensors
        )
        stage = 1 if ctx.causal else 0
        assert do.is_contiguous()
        assert q.stride() == k.stride() == v.stride() == o.stride() == do.stride()
        dq = torch.empty_like(q)
        dk = torch.empty_like(k)
        dv = torch.empty_like(v)
        BATCH, N_HEAD, N_CTX = q.shape[:3]
        PRE_BLOCK = 128
        NUM_WARPS, NUM_STAGES = 4, 5
        BLOCK_M1, BLOCK_N1, BLOCK_M2, BLOCK_N2 = 32, 128, 128, 32
        BLK_SLICE_FACTOR = 2
        RCP_LN2 = 1.4426950408889634  # = 1.0 / ln(2)
        arg_k = k
        arg_k = arg_k * (ctx.sm_scale * RCP_LN2)
        PRE_BLOCK = 128
        assert N_CTX % PRE_BLOCK == 0
        pre_grid = (N_CTX // PRE_BLOCK, BATCH * N_HEAD)
        delta = torch.empty_like(M)
        _attn_bwd_preprocess[pre_grid](
            o,
            do,  #
            delta,  #
            BATCH,
            N_HEAD,
            N_CTX,  #
            BLOCK_M=PRE_BLOCK,  # type: ignore
            HEAD_DIM=ctx.HEAD_DIM,  #
        )
        grid = (N_CTX // BLOCK_N1, 1, BATCH * N_HEAD)
        _attn_bwd[grid](
            q,
            arg_k,
            v,
            ctx.sm_scale,
            maskMat,
            blkBinaryMask,
            blkBinaryMask_T,  #
            num_ones,
            offset_ones,
            num_ones.stride(0),  #
            num_ones_T,
            offset_ones_T,
            num_ones_T.stride(0),  #
            do,
            dq,
            dk,
            dv,  #
            M,
            delta,  #
            q.stride(0),
            q.stride(1),
            q.stride(2),
            q.stride(3),  #
            maskMat.stride(0),
            maskMat.stride(1),  #
            blkBinaryMask.stride(0),
            blkBinaryMask.stride(1),  #
            blkBinaryMask_T.stride(0),
            blkBinaryMask_T.stride(1),  #
            N_HEAD,
            N_CTX,  #
            IS_DENSEBLK=stage,  #
            BLOCK_M1=BLOCK_M1,
            BLOCK_N1=BLOCK_N1,  #
            BLOCK_M2=BLOCK_M2,
            BLOCK_N2=BLOCK_N2,  #
            BLK_SLICE_FACTOR=BLK_SLICE_FACTOR,  #
            HEAD_DIM=ctx.HEAD_DIM,  #
            num_warps=NUM_WARPS,  # type: ignore
            num_stages=NUM_STAGES,  # type: ignore
        )

        return dq, dk, dv, None, None, None, None, None, None, None, None, None


attention_dense_binaryBlkMat = _attention_dense_binaryBlkMat.apply


@triton.jit
def _attn_base_fwd_inner(
    acc,
    l_i,
    m_i,
    q,  #
    K_block_ptr,
    V_block_ptr,
    Mask_block_ptr,  #
    blkBinaryMask_ptr,
    start_m,
    qk_scale,  #
    blkBinaryMask_stride_d1,
    blkBinaryMask_stride_d2,  # Strides for BLK_ptr
    BLOCK_M: tl.constexpr,
    HEAD_DIM: tl.constexpr,
    BLOCK_N: tl.constexpr,  #
    N_CTX: tl.constexpr,
    fp8_v: tl.constexpr,
):
    lo, hi = 0, N_CTX
    K_block_ptr = tl.advance(K_block_ptr, (0, lo))
    V_block_ptr = tl.advance(V_block_ptr, (lo, 0))
    Mask_block_ptr = tl.advance(Mask_block_ptr, (0, lo))
    # loop over k, v and update accumulator
    for start_n in range(lo, hi, BLOCK_N):
        cur_val = start_n // BLOCK_N
        start_n = tl.multiple_of(start_n, BLOCK_N)

        binaryBlk_ptr = start_m * blkBinaryMask_stride_d1 + cur_val * blkBinaryMask_stride_d2
        binaryBlk_val = tl.load(blkBinaryMask_ptr + binaryBlk_ptr)
        if binaryBlk_val:
            # -- compute qk ----
            k = tl.load(K_block_ptr)
            qk = tl.dot(q, k)

            ## Added for random masking case
            mask = tl.load(Mask_block_ptr)
            qk = qk * qk_scale + tl.where(mask, 0, -1.0e6)
            ## Move later to a separate STAGE
            m_ij = tl.maximum(m_i, tl.max(qk, 1))
            # qk = qk * qk_scale - m_ij[:, None]
            qk -= m_ij[:, None]
            ##### Separate block as I might move it for casual cases ####

            p = tl.math.exp2(qk)
            l_ij = tl.sum(p, 1)
            # -- update m_i and l_i
            alpha = tl.math.exp2(m_i - m_ij)
            l_i = l_i * alpha + l_ij
            # -- update output accumulator --
            acc = acc * alpha[:, None]
            # update acc
            v = tl.load(V_block_ptr)
            p = p.to(tl.float8e5) if fp8_v else p.to(q.type.element_ty)
            acc = tl.dot(p, v, acc)
            # update m_i and l_i
            m_i = m_ij
        V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0))
        K_block_ptr = tl.advance(K_block_ptr, (0, BLOCK_N))
        Mask_block_ptr = tl.advance(Mask_block_ptr, (0, BLOCK_N))
        # blkBinaryMask_ptr = tl.advance(blkBinaryMask_ptr, (0, 1))
    return acc, l_i, m_i


@triton.jit
def _attn_base_fwd(
    Q,
    K,
    V,
    sm_scale,
    M,
    Out,
    maskMat,
    blkBinaryMask,  #
    stride_qz,
    stride_qh,
    stride_qm,
    stride_qk,  #
    stride_kz,
    stride_kh,
    stride_kn,
    stride_kk,  #
    stride_vz,
    stride_vh,
    stride_vk,
    stride_vn,  #
    stride_oz,
    stride_oh,
    stride_om,
    stride_on,  #
    maskMat_stride_d1,
    maskMat_stride_d2,  #
    blkBinaryMask_stride_d1,
    blkBinaryMask_stride_d2,  #
    Z,
    H,
    N_CTX,  #
    HEAD_DIM: tl.constexpr,  #
    BLOCK_M: tl.constexpr,  #
    BLOCK_N: tl.constexpr,  #
):
    tl.static_assert(BLOCK_N <= HEAD_DIM)
    start_m = tl.program_id(0)
    off_hz = tl.program_id(1)
    off_z = off_hz // H
    off_h = off_hz % H
    qvk_offset = off_z.to(tl.int64) * stride_qz + off_h.to(tl.int64) * stride_qh

    # block pointers
    Q_block_ptr = tl.make_block_ptr(
        base=Q + qvk_offset,
        shape=(N_CTX, HEAD_DIM),
        strides=(stride_qm, stride_qk),
        offsets=(start_m * BLOCK_M, 0),
        block_shape=(BLOCK_M, HEAD_DIM),
        order=(1, 0),
    )
    v_order: tl.constexpr = (0, 1) if V.dtype.element_ty == tl.float8e5 else (1, 0)  # type: ignore
    V_block_ptr = tl.make_block_ptr(
        base=V + qvk_offset,
        shape=(N_CTX, HEAD_DIM),
        strides=(stride_vk, stride_vn),
        offsets=(0, 0),
        block_shape=(BLOCK_N, HEAD_DIM),
        order=v_order,
    )
    K_block_ptr = tl.make_block_ptr(
        base=K + qvk_offset,
        shape=(HEAD_DIM, N_CTX),
        strides=(stride_kk, stride_kn),
        offsets=(0, 0),
        block_shape=(HEAD_DIM, BLOCK_N),
        order=(0, 1),
    )
    O_block_ptr = tl.make_block_ptr(
        base=Out + qvk_offset,
        shape=(N_CTX, HEAD_DIM),
        strides=(stride_om, stride_on),
        offsets=(start_m * BLOCK_M, 0),
        block_shape=(BLOCK_M, HEAD_DIM),
        order=(1, 0),
    )
    Mask_block_ptr = tl.make_block_ptr(
        base=maskMat,
        shape=(N_CTX, N_CTX),
        strides=(maskMat_stride_d1, maskMat_stride_d2),
        offsets=(start_m * BLOCK_M, 0),
        block_shape=(BLOCK_M, BLOCK_N),
        order=(1, 0),
    )
    # blkBinaryMask_ptr = tl.make_block_ptr(
    #     base=blkBinaryMask,
    #     shape=(N_CTX // BLOCK_M, N_CTX // BLOCK_N),
    #     strides=(blkBinaryMask_stride_d1, blkBinaryMask_stride_d2),
    #     offsets=(start_m, 0),
    #     block_shape=(1, 1),
    #     order=(1, 0),
    # )
    # initialize offsets
    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
    # offs_n = tl.arange(0, BLOCK_N)

    # initialize pointer to m and l
    m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
    l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0
    acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32)
    # load scales
    qk_scale = sm_scale
    qk_scale *= 1.44269504  # 1/log(2)
    # load q: it will stay in SRAM throughout
    q = tl.load(Q_block_ptr)

    acc, l_i, m_i = _attn_base_fwd_inner(
        acc,
        l_i,
        m_i,
        q,
        K_block_ptr,
        V_block_ptr,
        Mask_block_ptr,  #
        blkBinaryMask,
        start_m,
        qk_scale,  #
        blkBinaryMask_stride_d1,
        blkBinaryMask_stride_d2,  #
        BLOCK_M,
        HEAD_DIM,
        BLOCK_N,  #
        N_CTX,
        V.dtype.element_ty == tl.float8e5,  #
    )  # type: ignore

    # epilogue
    m_i += tl.math.log2(l_i)
    acc = acc / l_i[:, None]
    m_ptrs = M + off_hz * N_CTX + offs_m
    tl.store(m_ptrs, m_i)
    tl.store(O_block_ptr, acc.to(Out.type.element_ty))


@triton.jit
def _attn_base_bwd_preprocess(
    O,
    DO,  #
    Delta,  #
    Z,
    H,
    N_CTX,  #
    BLOCK_M: tl.constexpr,
    HEAD_DIM: tl.constexpr,  #
):
    off_m = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
    off_hz = tl.program_id(1)
    off_n = tl.arange(0, HEAD_DIM)
    # load
    o = tl.load(O + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :])
    do = tl.load(DO + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :]).to(tl.float32)
    delta = tl.sum(o * do, axis=1)
    # write-back
    tl.store(Delta + off_hz * N_CTX + off_m, delta)


# The main inner-loop logic for computing dK and dV.
@triton.jit
def _attn_base_bwd_dkdv(
    dk,
    dv,  #
    Q,
    k,
    v,
    sm_scale,
    maskMat,
    blkBinaryMask,  #
    DO,  #
    M,
    D,  #
    # shared by Q/K/V/DO.
    stride_tok,
    stride_d,  #
    stride_m_d1,
    stride_m_d2,  #
    blkBinaryMask_stride_d1,
    blkBinaryMask_stride_d2,  #
    H,
    N_CTX,
    BLOCK_M1: tl.constexpr,  #
    BLOCK_N1: tl.constexpr,  #
    HEAD_DIM: tl.constexpr,  #
    # Filled in by the wrapper.
    start_n,
    start_m,
    num_steps,  #
    MASK: tl.constexpr,
):
    offs_m = start_m + tl.arange(0, BLOCK_M1)
    offs_n = start_n + tl.arange(0, BLOCK_N1)
    offs_k = tl.arange(0, HEAD_DIM)
    qT_ptrs = Q + offs_m[None, :] * stride_tok + offs_k[:, None] * stride_d
    do_ptrs = DO + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
    maskMat_T_ptrs = maskMat + offs_m[None, :] * stride_m_d1 + offs_n[:, None] * stride_m_d2

    # BLOCK_N1 must be a multiple of BLOCK_M1, otherwise the code wouldn't work.
    tl.static_assert(BLOCK_N1 % BLOCK_M1 == 0)
    curr_m = start_m
    step_m = BLOCK_M1
    curr_n = start_n // BLOCK_N1
    for blk_idx in range(num_steps):
        blkBinaryMask_ptr = curr_n * blkBinaryMask_stride_d1 + blk_idx * blkBinaryMask_stride_d2
        binaryBlk_val = tl.load(blkBinaryMask + blkBinaryMask_ptr)
        if binaryBlk_val:
            qT = tl.load(qT_ptrs)
            # Load m before computing qk to reduce pipeline stall.
            offs_m = curr_m + tl.arange(0, BLOCK_M1)
            m = tl.load(M + offs_m)
            qkT = tl.dot(k, qT)
            pT = tl.math.exp2(qkT - m[None, :])
            # Autoregressive masking.
            if MASK:
                mask = tl.load(maskMat_T_ptrs)
                pT = tl.where(mask, pT, 0.0)
            do = tl.load(do_ptrs)
            # Compute dV.
            ppT = pT
            ppT = ppT.to(Q.type.element_ty)
            dv += tl.dot(ppT, do)
            # D (= delta) is pre-divided by ds_scale.
            Di = tl.load(D + offs_m)
            # Compute dP and dS.
            dpT = tl.dot(v, tl.trans(do)).to(tl.float32)
            dsT = pT * (dpT - Di[None, :])
            dsT = dsT.to(Q.type.element_ty)
            dk += tl.dot(dsT, tl.trans(qT))
        # Increment pointers.
        curr_m += step_m
        qT_ptrs += step_m * stride_tok
        do_ptrs += step_m * stride_tok
        maskMat_T_ptrs += step_m * stride_m_d1
    return dk, dv


# the main inner-loop logic for computing dQ
@triton.jit
def _attn_base_bwd_dq(
    dq,
    q,
    K,
    V,
    maskMat,
    blkBinaryMask,  #
    do,
    m,
    D,
    # shared by Q/K/V/DO.
    stride_tok,
    stride_d,  #
    stride_m_d1,
    stride_m_d2,  #
    blkBinaryMask_stride_d1,
    blkBinaryMask_stride_d2,  #
    H,
    N_CTX,  #
    BLOCK_M2: tl.constexpr,  #
    BLOCK_N2: tl.constexpr,  #
    HEAD_DIM: tl.constexpr,
    # Filled in by the wrapper.
    start_m,
    start_n,
    num_steps,  #
    MASK: tl.constexpr,
):
    offs_m = start_m + tl.arange(0, BLOCK_M2)
    offs_n = start_n + tl.arange(0, BLOCK_N2)
    offs_k = tl.arange(0, HEAD_DIM)
    kT_ptrs = K + offs_n[None, :] * stride_tok + offs_k[:, None] * stride_d
    vT_ptrs = V + offs_n[None, :] * stride_tok + offs_k[:, None] * stride_d
    maskMat_ptrs = maskMat + offs_m[:, None] * stride_m_d1 + offs_n[None, :] * stride_m_d2
    # D (= delta) is pre-divided by ds_scale.
    Di = tl.load(D + offs_m)
    # BLOCK_M2 must be a multiple of BLOCK_N2, otherwise the code wouldn't work.
    tl.static_assert(BLOCK_M2 % BLOCK_N2 == 0)
    curr_n = start_n
    step_n = BLOCK_N2
    curr_m = start_m // BLOCK_M2
    for blk_idx in range(num_steps):
        binaryBlk_ptr = curr_m * blkBinaryMask_stride_d1 + blk_idx * blkBinaryMask_stride_d2
        binaryBlk_val = tl.load(blkBinaryMask + binaryBlk_ptr)
        # print('blk_idx', blk_idx)
        if binaryBlk_val:
            kT = tl.load(kT_ptrs)
            vT = tl.load(vT_ptrs)
            qk = tl.dot(q, kT)
            p = tl.math.exp2(qk - m)
            # Autoregressive masking.
            if MASK:
                mask = tl.load(maskMat_ptrs)
                # offs_n = curr_n + tl.arange(0, BLOCK_N2)
                # mask = (offs_m[:, None] >= offs_n[None, :])
                p = tl.where(mask, p, 0.0)
            # Compute dP and dS.
            dp = tl.dot(do, vT).to(tl.float32)
            ds = p * (dp - Di[:, None])
            ds = ds.to(q.type.element_ty)
            # Compute dQ.
            # NOTE: We need to de-scale dq in the end, because kT was pre-scaled.
            dq += tl.dot(ds, tl.trans(kT))
        # Increment pointers.
        curr_n += step_n
        kT_ptrs += step_n * stride_tok
        vT_ptrs += step_n * stride_tok
        maskMat_ptrs += step_n * stride_m_d2

    return dq


@triton.jit
def _attn_base_bwd(
    Q,
    K,
    V,
    sm_scale,
    maskMat,
    blkBinaryMask,
    blkBinaryMask_T,  #
    DO,  #
    DQ,
    DK,
    DV,  #
    M,
    D,
    # shared by Q/K/V/DO.
    stride_z,
    stride_h,
    stride_tok,
    stride_d,  #
    stride_m_d1,
    stride_m_d2,  #
    blkBinaryMask_stride_d1,
    blkBinaryMask_stride_d2,  #
    blkBinaryMask_T_stride_d1,
    blkBinaryMask_T_stride_d2,  #
    H,
    N_CTX,  #
    BLOCK_M1: tl.constexpr,  #
    BLOCK_N1: tl.constexpr,  #
    BLOCK_M2: tl.constexpr,  #
    BLOCK_N2: tl.constexpr,  #
    BLK_SLICE_FACTOR: tl.constexpr,  #
    HEAD_DIM: tl.constexpr,
):
    LN2: tl.constexpr = 0.6931471824645996  # type: ignore # = ln(2)

    bhid = tl.program_id(2)
    off_chz = (bhid * N_CTX).to(tl.int64)
    adj = (stride_h * (bhid % H) + stride_z * (bhid // H)).to(tl.int64)
    pid = tl.program_id(0)

    # offset pointers for batch/head
    Q += adj
    K += adj
    V += adj
    DO += adj
    DQ += adj
    DK += adj
    DV += adj
    M += off_chz
    D += off_chz
    ## NO MaskMat update till here

    # load scales
    offs_k = tl.arange(0, HEAD_DIM)

    start_n = pid * BLOCK_N1
    # start_m = start_n
    start_m = 0

    # MASK_BLOCK_M1: tl.constexpr = BLOCK_M1 // BLK_SLICE_FACTOR
    offs_n = start_n + tl.arange(0, BLOCK_N1)

    dv = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)
    dk = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32)

    # load K and V: they stay in SRAM throughout the inner loop.
    k = tl.load(K + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)
    v = tl.load(V + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d)

    # num_steps = BLOCK_N1 // MASK_BLOCK_M1

    # dk, dv = _attn_bwd_dkdv(dk, dv,  #
    #                         Q, k, v, sm_scale,  #
    #                         DO,  #
    #                         M, D,  #
    #                         stride_tok, stride_d,  #
    #                         H, N_CTX,  #
    #                         MASK_BLOCK_M1, BLOCK_N1, HEAD_DIM,  #
    #                         start_n, start_m, num_steps,  #
    #                         MASK=True  #
    #                         )

    # start_m += num_steps * MASK_BLOCK_M1
    num_steps = (N_CTX) // BLOCK_M1

    # Compute dK and dV for non-masked blocks.
    dk, dv = _attn_base_bwd_dkdv(  #
        dk,
        dv,  #
        Q,
        k,
        v,
        sm_scale,
        maskMat,
        blkBinaryMask_T,  #
        DO,  #
        M,
        D,  #
        stride_tok,
        stride_d,  #
        stride_m_d1,
        stride_m_d2,
        blkBinaryMask_T_stride_d1,
        blkBinaryMask_T_stride_d2,  #
        H,
        N_CTX,  #
        BLOCK_M1,
        BLOCK_N1,
        HEAD_DIM,  #
        start_n,
        start_m,
        num_steps,  #
        MASK=True,  #
    )  # type: ignore

    dv_ptrs = DV + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
    tl.store(dv_ptrs, dv)

    # Write back dK.
    dk *= sm_scale
    dk_ptrs = DK + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d
    tl.store(dk_ptrs, dk)

    # THIS BLOCK DOES DQ:
    start_m = pid * BLOCK_M2
    end_n = N_CTX

    # MASK_BLOCK_N2: tl.constexpr = BLOCK_N2 // BLK_SLICE_FACTOR
    offs_m = start_m + tl.arange(0, BLOCK_M2)

    q = tl.load(Q + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d)
    dq = tl.zeros([BLOCK_M2, HEAD_DIM], dtype=tl.float32)
    do = tl.load(DO + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d)

    m = tl.load(M + offs_m)
    m = m[:, None]

    # Compute dQ for masked (diagonal) blocks.
    # NOTE: This code scans each row of QK^T backward (from right to left,
    # but inside each call to _attn_bwd_dq, from left to right), but that's
    # not due to anything important.  I just wanted to reuse the loop
    # structure for dK & dV above as much as possible.
    # num_steps = BLOCK_M2 // MASK_BLOCK_N2
    # dq = _attn_bwd_dq(dq, q, K, V,  #
    #                   do, m, D,  #
    #                   stride_tok, stride_d,  #
    #                   H, N_CTX,  #
    #                   BLOCK_M2, MASK_BLOCK_N2, HEAD_DIM,  #
    #                   start_m, end_n - num_steps * MASK_BLOCK_N2, num_steps,  #
    #                   MASK=True  #
    #                   )
    # end_n -= num_steps * MASK_BLOCK_N2
    # # stage 2

    num_steps = end_n // BLOCK_N2
    dq = _attn_base_bwd_dq(
        dq,
        q,
        K,
        V,
        maskMat,
        blkBinaryMask,  #
        do,
        m,
        D,  #
        stride_tok,
        stride_d,  #
        stride_m_d1,
        stride_m_d2,  #
        blkBinaryMask_stride_d1,
        blkBinaryMask_stride_d2,  #
        H,
        N_CTX,  #
        BLOCK_M2,
        BLOCK_N2,
        HEAD_DIM,  #
        start_m,
        end_n - num_steps * BLOCK_N2,
        num_steps,  #
        MASK=True,  #
    )
    # Write back dQ.
    dq_ptrs = DQ + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d
    dq *= LN2
    tl.store(dq_ptrs, dq)


class _attention_base_binaryBlkMat(torch.autograd.Function):
    """
    This is a cleaned up version of the Binary Block Masked Attention.
    It doesn't contain the casual masking logic.
    """

    @staticmethod
    def forward(ctx, q, k, v, causal, sm_scale, maskMat, blkBinaryMask, blkBinaryMask_T):
        # shape constraints
        HEAD_DIM_Q, HEAD_DIM_K = q.shape[-1], k.shape[-1]
        # when v is in float8_e5m2 it is transposed.
        HEAD_DIM_V = v.shape[-1]
        assert HEAD_DIM_Q == HEAD_DIM_K and HEAD_DIM_K == HEAD_DIM_V
        assert HEAD_DIM_K in {16, 32, 64, 128, 256}
        o = torch.empty_like(q)

        extra_kern_args = {}
        # Tuning for AMD target
        if is_hip():
            waves_per_eu = 3 if HEAD_DIM_K <= 64 else 2
            extra_kern_args = {"waves_per_eu": waves_per_eu, "allow_flush_denorm": True}

        grid = lambda args: (triton.cdiv(q.shape[2], args["BLOCK_M"]), q.shape[0] * q.shape[1], 1)
        M = torch.empty((q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32)
        _attn_base_fwd[grid](
            q,
            k,
            v,
            sm_scale,
            M,
            o,
            maskMat,
            blkBinaryMask,  #
            q.stride(0),
            q.stride(1),
            q.stride(2),
            q.stride(3),  #
            k.stride(0),
            k.stride(1),
            k.stride(2),
            k.stride(3),  #
            v.stride(0),
            v.stride(1),
            v.stride(2),
            v.stride(3),  #
            o.stride(0),
            o.stride(1),
            o.stride(2),
            o.stride(3),  #
            maskMat.stride(0),
            maskMat.stride(1),  #
            blkBinaryMask.stride(0),
            blkBinaryMask.stride(1),  #
            q.shape[0],
            q.shape[1],  #
            N_CTX=q.shape[2],  #
            HEAD_DIM=HEAD_DIM_K,  #
            BLOCK_M=128,  # type: ignore
            BLOCK_N=32,  # type: ignore
            **extra_kern_args,
        )

        ctx.save_for_backward(q, k, v, o, M, maskMat, blkBinaryMask, blkBinaryMask_T)
        ctx.grid = grid
        ctx.sm_scale = sm_scale
        ctx.HEAD_DIM = HEAD_DIM_K
        ctx.causal = causal
        return o

    @staticmethod
    def backward(ctx, do):
        q, k, v, o, M, maskMat, blkBinaryMask, blkBinaryMask_T = ctx.saved_tensors
        assert do.is_contiguous()
        assert q.stride() == k.stride() == v.stride() == o.stride() == do.stride()
        dq = torch.empty_like(q)
        dk = torch.empty_like(k)
        dv = torch.empty_like(v)
        BATCH, N_HEAD, N_CTX = q.shape[:3]
        PRE_BLOCK = 128
        NUM_WARPS, NUM_STAGES = 4, 5
        BLOCK_M1, BLOCK_N1, BLOCK_M2, BLOCK_N2 = 32, 128, 128, 32
        BLK_SLICE_FACTOR = 2
        RCP_LN2 = 1.4426950408889634  # = 1.0 / ln(2)
        arg_k = k
        arg_k = arg_k * (ctx.sm_scale * RCP_LN2)
        PRE_BLOCK = 128
        assert N_CTX % PRE_BLOCK == 0
        pre_grid = (N_CTX // PRE_BLOCK, BATCH * N_HEAD)
        delta = torch.empty_like(M)
        _attn_base_bwd_preprocess[pre_grid](
            o,
            do,  #
            delta,  #
            BATCH,
            N_HEAD,
            N_CTX,  #
            BLOCK_M=PRE_BLOCK,  # type: ignore
            HEAD_DIM=ctx.HEAD_DIM,  #
        )
        grid = (N_CTX // BLOCK_N1, 1, BATCH * N_HEAD)
        _attn_base_bwd[grid](
            q,
            arg_k,
            v,
            ctx.sm_scale,
            maskMat,
            blkBinaryMask,
            blkBinaryMask_T,
            do,
            dq,
            dk,
            dv,  #
            M,
            delta,  #
            q.stride(0),
            q.stride(1),
            q.stride(2),
            q.stride(3),  #
            maskMat.stride(0),
            maskMat.stride(1),  #
            blkBinaryMask.stride(0),
            blkBinaryMask.stride(1),  #
            blkBinaryMask_T.stride(0),
            blkBinaryMask_T.stride(1),  #
            N_HEAD,
            N_CTX,  #
            BLOCK_M1=BLOCK_M1,
            BLOCK_N1=BLOCK_N1,  #
            BLOCK_M2=BLOCK_M2,
            BLOCK_N2=BLOCK_N2,  #
            BLK_SLICE_FACTOR=BLK_SLICE_FACTOR,  #
            HEAD_DIM=ctx.HEAD_DIM,  #
            num_warps=NUM_WARPS,  # type: ignore
            num_stages=NUM_STAGES,  # type: ignore
        )  # type: ignore

        return dq, dk, dv, None, None, None, None, None


attention_binaryBlkMat = _attention_base_binaryBlkMat.apply


@triton.jit
def calculate_sum_matrix(
    input_mat,
    output_sum_matrix,
    output_binblk_matrix,
    input_stride_0,
    input_stride_1,
    output_stride_0,
    output_stride_1,
    BLKSIZE_I: tl.constexpr,
    BLKSIZE_J: tl.constexpr,
):
    row = tl.program_id(0)
    col = tl.program_id(1)

    # Get the relevant block from the input matrix
    start_m = row * BLKSIZE_I
    start_n = col * BLKSIZE_J

    offs_m = start_m + tl.arange(0, BLKSIZE_I)
    offs_n = start_n + tl.arange(0, BLKSIZE_J)

    input_ptrs = input_mat + offs_m[:, None] * input_stride_0 + offs_n[None, :] * input_stride_1
    input_block = tl.load(input_ptrs)

    # Calculate the sum of the block
    sum_block = tl.sum(input_block)

    # print("row_offset", row_offset.shape[1])
    # Store the sum in the output sum matrix
    output_sum_ptr = output_sum_matrix + row * output_stride_0 + col * output_stride_1
    tl.store(output_sum_ptr, sum_block)

    # Store the binary block in the output binBlk matrix
    output_binblk_ptr = output_binblk_matrix + row * output_stride_0 + col * output_stride_1
    tl.store(output_binblk_ptr, (sum_block > 0).to(output_binblk_matrix.type.element_ty))


def return_sum_matrix(test_mat, BLKSIZE_I: tl.constexpr = 128, BLKSIZE_J: tl.constexpr = 32):  # type: ignore
    sum_matrix = torch.zeros(
        (test_mat.size(0) // BLKSIZE_I, test_mat.size(1) // BLKSIZE_J), device=test_mat.device, dtype=torch.float32
    )
    binblk_matrix = torch.zeros_like(sum_matrix)
    grid = (sum_matrix.size(0), sum_matrix.size(1))

    calculate_sum_matrix[grid](
        test_mat,
        sum_matrix,
        binblk_matrix,
        test_mat.stride(0),
        test_mat.stride(1),
        sum_matrix.stride(0),
        sum_matrix.stride(1),
        BLKSIZE_I,
        BLKSIZE_J,
    )
    return sum_matrix, binblk_matrix


def return_sum_matrix_v2(testMat, BLKSIZE_I=128, BLKSIZE_J=32):
    # Get the dimensions of the input matrix
    nrows, ncols = testMat.shape

    # Reshape the matrix to create blocks
    reshaped = testMat.view(nrows // BLKSIZE_I, BLKSIZE_I, -1, BLKSIZE_J)

    # Sum the elements in each block
    block_sums = reshaped.sum(dim=(1, 3))

    return block_sums


@triton.jit
def calculate_ones_and_offset(
    input_sum_matrix,
    output_ones_matrix,
    output_offset_matrix,
    input_stride_0,
    input_stride_1,
    output_stride_0,
    output_stride_1,
    SUM_CHECK: tl.constexpr,
    COL_SIZE: tl.constexpr,
):
    row = tl.program_id(0)

    # Get the relevant block from the input matrix
    offs_n = tl.arange(0, COL_SIZE)
    input_ptrs = input_sum_matrix + row * input_stride_0 + offs_n[None, :] * input_stride_1
    input_block = tl.load(input_ptrs)

    # Check complete one block
    cond_inp_blk = tl.where(input_block == SUM_CHECK, 1, 0)

    # Store total complete ones in the output ones matrix
    tot_ones = tl.sum(cond_inp_blk)
    output_ones_ptr = output_ones_matrix + row * output_stride_0
    tl.store(output_ones_ptr, tot_ones)

    # Store the offset in the output offset matrix
    first_one = tl.argmax(cond_inp_blk, axis=1)
    off_ptr = tl.arange(0, 1)
    output_offset_ptr = output_offset_matrix + row * output_stride_0 + off_ptr * output_stride_1
    tl.store(output_offset_ptr, first_one)


def return_ones_and_offset(sum_matrix, SUM_CHECK):
    ones_mat = torch.zeros((sum_matrix.size(0), 1), dtype=sum_matrix.dtype, device=sum_matrix.device)
    offset_mat = torch.zeros_like(ones_mat)
    grid = (sum_matrix.size(0), 1)

    COL_SIZE = sum_matrix.size(1)
    calculate_ones_and_offset[grid](
        sum_matrix,
        ones_mat,
        offset_mat,
        sum_matrix.stride(0),
        sum_matrix.stride(1),
        ones_mat.stride(0),
        ones_mat.stride(1),
        SUM_CHECK,
        COL_SIZE,
    )
    return ones_mat[:, 0], offset_mat[:, 0]


def return_binBlk_matrices(test_mat, BLKSIZE_I=128, BLKSIZE_J=32, is_dense=False):
    SUM_CHECK = BLKSIZE_I * BLKSIZE_J

    sum_matrix, binBlk_matrix = return_sum_matrix(test_mat, BLKSIZE_I, BLKSIZE_J)  # type: ignore
    if is_dense:
        ones_matrix, offset_matrix = return_ones_and_offset(sum_matrix, SUM_CHECK)
        return binBlk_matrix, ones_matrix, offset_matrix
    return binBlk_matrix


def binBlk_wrapper(q, k, v, mask=None, causal=True, dense=True):
    _, S, _, hs = q.shape
    q = q.transpose(1, 2)  # (B, nh, S, hs)
    k = k.transpose(1, 2)
    v = v.transpose(1, 2)

    if mask is None:
        mask = torch.tril(torch.ones((S, S), device=q.device, dtype=q.dtype))
    maskMat_blk, num_ones, offset = return_binBlk_matrices(mask, is_dense=True)
    maskMat_blk_T, num_ones_T, offset_T = return_binBlk_matrices(mask.T, is_dense=True)

    sm_scale = 1 / math.sqrt(hs)

    if dense:
        y = attention_dense_binaryBlkMat(
            q,
            k,
            v,
            causal,
            sm_scale,
            mask,
            maskMat_blk,
            maskMat_blk_T,
            num_ones,
            offset,
            num_ones_T,
            offset_T,
        )
    else:
        y = attention_binaryBlkMat(
            q,
            k,
            v,
            causal,
            sm_scale,
            mask,
            maskMat_blk,
            maskMat_blk_T,
            num_ones,
            offset,
            num_ones_T,
            offset_T,
        )

    return y.transpose(1, 2)  # type: ignore
